{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### 关于 `Pytorch` 中的几个重要类 和代码格式的介绍",
   "id": "3216b454c1be63b2"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "#### 重要类",
   "id": "efa56a2750760ba7"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "`MyModel` 继承 `nn.Module`",
   "id": "179457982a0d9532"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class MyModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Linear(10, 2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc(x)\n"
   ],
   "id": "7e7cbff72c26ff5b"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "torch.utils.data.Dataset 数据集类\n",
    "| 方法                         | 作用                                      |\n",
    "| -------------------------- | --------------------------------------- |\n",
    "| `__len__(self)`            | 返回数据集中样本的总数                             |\n",
    "| `__getitem__(self, index)` | 返回指定索引 `index` 的样本（通常为 `(data, label)`） |\n"
   ],
   "id": "3ac96f8fd0a7dbf8"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "from torch.utils.data import Dataset\n",
    "from PIL import Image\n",
    "import os\n",
    "\n",
    "class MyImageDataset(Dataset):\n",
    "    def __init__(self, img_dir, transform=None):\n",
    "        self.img_dir = img_dir\n",
    "        self.transform = transform\n",
    "        self.img_files = os.listdir(img_dir)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_files)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_name = self.img_files[idx]\n",
    "        img_path = os.path.join(self.img_dir, img_name)\n",
    "        image = Image.open(img_path).convert(\"RGB\")\n",
    "        label = int(img_name.split('_')[0])  # 假设文件名里带标签\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        return image, label\n"
   ],
   "id": "19ed9639bef95925"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "torch.utils.data.DataLoader —— 数据加载器\n",
    "\n",
    "DataLoader 是数据集的加载器，用于：\n",
    "\n",
    "- 批量读取数据（batch）\n",
    "- 打乱顺序（shuffle）\n",
    "- 并行读取（num_workers）\n",
    "- 自动将数据打包成 mini-batch\n",
    "| 参数            | 说明                |\n",
    "| ------------- | ----------------- |\n",
    "| `dataset`     | 传入一个 `Dataset` 对象 |\n",
    "| `batch_size`  | 每批样本数             |\n",
    "| `shuffle`     | 是否在每个 epoch 打乱数据  |\n",
    "| `num_workers` | 读取数据的线程数（>0 可加速）  |\n",
    "| `drop_last`   | 若最后一个 batch 不足则丢弃 |\n",
    "| `collate_fn`  | 自定义合并 batch 的函数   |\n"
   ],
   "id": "cbfc85c3922a5bbe"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# 实例化自定义数据集\n",
    "dataset = MyImageDataset(img_dir='data/images', transform=None)\n",
    "\n",
    "# 用 DataLoader 封装\n",
    "dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)\n",
    "\n",
    "# 使用示例\n",
    "for batch_idx, (images, labels) in enumerate(dataloader):\n",
    "    print(batch_idx, images.shape, labels)\n"
   ],
   "id": "9663a55a939e98ee"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "#### 基本写法",
   "id": "69e010e8bd8129a2"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "```python\n",
    "# 导入模块\n",
    "\n",
    "# ========== 1. 数据加载 ==========\n",
    "\n",
    "# 导入数据集\n",
    "\n",
    "# 使用DataLoader 导入\n",
    "train_loader = DataLoader(xxx)\n",
    "test_loader  = DataLoader(xxx)\n",
    "\n",
    "# ========== 2. 定义模型 ==========\n",
    "\n",
    "class CNNModel(nn.Module):  # 继承\n",
    "    def __init__(self):\n",
    "        super(CNNModel, self).__init__()\n",
    "        # 定义相关层\n",
    "\n",
    "    def forward(self, x):  # 定义反向传播\n",
    "        output = xxx\n",
    "        return output\n",
    "\n",
    "# ========== 3. 初始化 ==========\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = CNNModel().to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)  # 选择优化器\n",
    "\n",
    "# ========== 4. 训练与验证 ==========\n",
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()  # 梯度清零\n",
    "        output = model(data)  # 输出\n",
    "        loss = F.nll_loss(output, target)  # 计算误差\n",
    "        loss.backward()  # 反向传播\n",
    "        optimizer.step()  # 更新\n",
    "        total_loss += loss.item()\n",
    "    print(f\"Epoch {epoch}: Train loss = {total_loss / len(train_loader):.4f}\")\n",
    "\n",
    "def test(model, device, test_loader):  # 验证模块\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += F.nll_loss(output, target, reduction='sum').item()\n",
    "            pred = output.argmax(dim=1)\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    acc = correct / len(test_loader.dataset)\n",
    "    print(f\"Test set: Average loss = {test_loss:.4f}, Accuracy = {acc*100:.2f}%\")\n",
    "    return acc\n",
    "\n",
    "# ========== 5. 主训练循环 ==========\n",
    "epochs = 5\n",
    "for epoch in range(1, epochs + 1):\n",
    "    train(model, device, train_loader, optimizer, epoch)\n",
    "    test(model, device, test_loader)\n",
    "\n",
    "# ========== 6. 保存模型 ==========\n",
    "torch.save(model.state_dict(), \"mnist_cnn.pt\")\n",
    "print(\"✅ 模型已保存为 mnist_cnn.pt\")\n",
    "\n",
    "```\n",
    "下例为基于Minist数据集的手写数字识别"
   ],
   "id": "ec8c47ca58c4e7aa"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# mnist_pytorch.py\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "# ========== 1. 数据加载 ==========\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),                     # 转为 [0,1]\n",
    "    transforms.Normalize((0.1307,), (0.3081,)) # 标准化 MNIST 均值/方差\n",
    "])\n",
    "\n",
    "train_dataset = datasets.MNIST(root=\"./data\", train=True, transform=transform, download=True)\n",
    "test_dataset  = datasets.MNIST(root=\"./data\", train=False, transform=transform, download=True)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "test_loader  = DataLoader(test_dataset, batch_size=1000, shuffle=False)\n",
    "\n",
    "# ========== 2. 定义 CNN 模型 ==========\n",
    "class CNNModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CNNModel, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, 1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3, 1)\n",
    "        self.dropout1 = nn.Dropout(0.25)\n",
    "        self.dropout2 = nn.Dropout(0.5)\n",
    "        self.fc1 = nn.Linear(9216, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.conv1(x))\n",
    "        x = F.relu(self.conv2(x))\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = self.dropout1(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.dropout2(x)\n",
    "        x = self.fc2(x)\n",
    "        output = F.log_softmax(x, dim=1)\n",
    "        return output\n",
    "\n",
    "# ========== 3. 初始化 ==========\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = CNNModel().to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# ========== 4. 训练与验证 ==========\n",
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "    print(f\"Epoch {epoch}: Train loss = {total_loss / len(train_loader):.4f}\")\n",
    "\n",
    "def test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += F.nll_loss(output, target, reduction='sum').item()\n",
    "            pred = output.argmax(dim=1)\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    acc = correct / len(test_loader.dataset)\n",
    "    print(f\"Test set: Average loss = {test_loss:.4f}, Accuracy = {acc*100:.2f}%\")\n",
    "    return acc\n",
    "\n",
    "# ========== 5. 主训练循环 ==========\n",
    "epochs = 5\n",
    "for epoch in range(1, epochs + 1):\n",
    "    train(model, device, train_loader, optimizer, epoch)\n",
    "    test(model, device, test_loader)\n",
    "\n",
    "# ========== 6. 保存模型 ==========\n",
    "torch.save(model.state_dict(), \"mnist_cnn.pt\")\n",
    "print(\"✅ 模型已保存为 mnist_cnn.pt\")\n",
    "\n",
    "# ========== 7. 单张图像预测 ==========\n",
    "def predict_image(img_path, model_path=\"mnist_cnn.pt\"):\n",
    "    model = CNNModel()\n",
    "    model.load_state_dict(torch.load(model_path, map_location='cpu'))\n",
    "    model.eval()\n",
    "\n",
    "    # 读入灰度图像\n",
    "    img = Image.open(img_path).convert('L')\n",
    "    img = img.resize((28, 28))\n",
    "    arr = np.array(img).astype(np.float32)\n",
    "\n",
    "    # 反转颜色（白底黑字 → 黑底白字）\n",
    "    if arr.mean() > 127:\n",
    "        arr = 255 - arr\n",
    "\n",
    "    arr = arr / 255.0\n",
    "    arr = (arr - 0.1307) / 0.3081  # 同样的标准化\n",
    "    tensor = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0)  # (1,1,28,28)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output = model(tensor)\n",
    "        pred = output.argmax(dim=1).item()\n",
    "        confidence = torch.exp(output[0, pred]).item()\n",
    "\n",
    "    plt.imshow(arr, cmap='gray')\n",
    "    plt.title(f\"Predicted: {pred} (conf={confidence:.3f})\")\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "    return pred, confidence\n",
    "\n",
    "# 示例（替换为你自己的图片路径）\n",
    "# predict_image(\"my_digit.png\")\n"
   ],
   "id": "3452a337f500b432"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
